from train_and_eval_utils import get_train_info, eval_from_pb, export_pb_file, eval_from_chekckpoint


def test_get_train_info():
    log = ""
    with open("/home/tang/Object_Detection_Cloud/nohup.out", 'r') as f:
        lines = f.readlines()
        for line in lines:
            log += line
    ret = get_train_info(log)
    assert ret[133949] == 0.3697


def test_eval_from_pb():
    pb_path = "/home/tang/workspace/quechao/total/frozen_inference_graph_244135_total.pb"
    images_dir = "/home/tang/workspace/quechao"
    label_map_path = "/home/tang/Object_Detection_Cloud/test/test_label_map.pbtxt"
    threshold = 0.5
    ret = eval_from_pb(pb_path, images_dir, label_map_path, threshold)
    print(ret)


def test_export_pb_file():
    checkpoint_prefix = "/media/tang/15F6-2313/model/train/model.ckpt-7512"
    config_path = "/media/tang/15F6-2313/model/train/pipeline.config"
    output_dir = "/media/tang/15F6-2313/model/output"
    export_pb_file(checkpoint_prefix, config_path, output_dir)


def test_eval_from_ckpt():
    eval_from_chekckpoint("/media/tang/15F6-2313/model/trai", "", "", "", 0.1)
